import numpy as np
from scipy.optimize import curve_fit
from matplotlib import pyplot as plt

AFILE = "all.out"
NFILE = "subset/i{}.out"
N = 44
COLOR1 = (.2, .6, .9)
COLOR2 = (.6, .1, .2)
S = 1000
s = 100
ticks = lambda n, m: zip(*[(10**i, str(10**i)) for i in range(n, m)])

# Zipf's law

# Get experimental values
with open(AFILE) as fp:
    counts = []
    for p in fp.read().split("\n")[0:-2]:
        word, count = p.split(" ")
        counts.append(int(count))
x = list(range(len(counts)))
y = sorted(counts, reverse=True)

# Find approximation to law
total = sum(counts)
xa = np.linspace(1, 70000, 70000)
zipf = lambda n, a, b, c: total * c / ((n + b)**a)

# Starting at 100
xs = x[s:]
ys = y[s:]
a, b, c = curve_fit(zipf, xs, ys)[0]
print("a: {} b: {} c: {}".format(a, b, c))
ysa = [zipf(n, a, b, c) for n in xa]

plt.plot(x, y, '.', color=COLOR1)
plt.plot(xa, ysa, color=COLOR2, linewidth=2)
plt.yscale('log', basex=10)
plt.xscale('log', basex=10)
plt.axvspan(1, s, facecolor='0.92', edgecolor='0.92')
plt.xticks(*ticks(0, 6))
plt.yticks(*ticks(0, 6))
plt.title("Resultat")
plt.ylabel("Ocurrencies")
plt.xlabel("N-essima paraula")
plt.savefig("zipf-100.png")
plt.clf()

# Starting at 1000
xS = x[S:]
yS = y[S:]
a, b, c = curve_fit(zipf, xS, yS)[0]
print("a: {} b: {} c: {}".format(a, b, c))
ySa = [zipf(n, a, b, c) for n in xS]

plt.plot(xS, yS, '.', color=COLOR1)
plt.plot(xS, ySa, color=COLOR2, linewidth=2)
plt.yscale('log', basex=10)
plt.xscale('log', basex=10)
plt.xticks(*ticks(3, 6))
plt.yticks(*ticks(0, 4))
plt.ylim([1, 10**3])
plt.title("Resultat")
plt.ylabel("Ocurrencies")
plt.xlabel("N-essima paraula")
plt.savefig("zipf-1000.png")
plt.clf()

# Heaps' law

# Get experimental values
x = []
y = []
for i in range(1, N):
    with open(NFILE.format(i)) as fp:
        counts = []
        for p in fp.read().split("\n")[0:-2]:
            word, count = p.split(" ")
            counts.append(int(count))
        x.append(sum(counts))
        y.append(len(counts))

# Find approximation to law
heaps = lambda n, k, b: k * (n**b)
k, b = curve_fit(heaps, x, y)[0]
print("k: {} b: {}".format(k, b))
xp = np.linspace(min(x), 800000, 100)
yp = [heaps(n, k, b) for n in xp]

plt.plot(x, y, '.', color=COLOR1)
plt.plot(xp, yp, color=COLOR2, linewidth=2)
plt.title("Resultat")
plt.xlabel("Paraules totals")
plt.ylabel("Paraules diferents")
plt.savefig("heaps.png")
